{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "<style>\n",
    ".cell-output-ipywidget-background {\n",
    "    background-color: transparent !important;\n",
    "}\n",
    ":root {\n",
    "    --jp-widgets-color: var(--vscode-editor-foreground);\n",
    "    --jp-widgets-font-size: var(--vscode-editor-font-size);\n",
    "}  \n",
    "</style>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping tuning as there is no suitable data. This can happen when all the trajectories in the same group have the same reward and thus no advantage to train on.\n",
      "Advanced step from 264 to 265 (no training occurred)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96f344d734164144b12b3480cefc39a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "gather:   0%|          | 0/1152 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping tuning as there is no suitable data. This can happen when all the trajectories in the same group have the same reward and thus no advantage to train on.\n",
      "Advanced step from 265 to 266 (no training occurred)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6bd3b26ae80d4d61b20a5335693d5b89",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "gather:   0%|          | 0/1152 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping tuning as there is no suitable data. This can happen when all the trajectories in the same group have the same reward and thus no advantage to train on.\n",
      "Advanced step from 266 to 267 (no training occurred)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "174a88b44d0b4f53a8910ba4ff463215",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "gather:   0%|          | 0/1152 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skipping tuning as there is no suitable data. This can happen when all the trajectories in the same group have the same reward and thus no advantage to train on.\n",
      "Advanced step from 267 to 268 (no training occurred)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89e8164861604905a115fe6f53d26c35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "gather:   0%|          | 0/1152 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import openai\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "import art\n",
    "from art.local import LocalBackend\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "backend = LocalBackend()\n",
    "model = art.TrainableModel(\n",
    "    name=\"009\",\n",
    "    project=\"yes-no-maybe\",\n",
    "    base_model=\"Qwen/Qwen2.5-7B-Instruct\",\n",
    "    # _internal_config=art.dev.InternalModelConfig(\n",
    "    #     _decouple_vllm_and_unsloth=True,\n",
    "    #     engine_args=art.dev.EngineArgs(gpu_memory_utilization=0.7),\n",
    "    # ),\n",
    ")\n",
    "await model.register(backend)\n",
    "\n",
    "\n",
    "async def rollout(client: openai.AsyncOpenAI, prompt: str) -> art.Trajectory:\n",
    "    messages: art.Messages = [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": prompt,\n",
    "        }\n",
    "    ]\n",
    "    chat_completion = await client.chat.completions.create(\n",
    "        messages=messages, model=model.name, max_tokens=100, timeout=100\n",
    "    )\n",
    "    choice = chat_completion.choices[0]\n",
    "    content = choice.message.content\n",
    "    assert isinstance(content, str)\n",
    "    if content == \"yes\":\n",
    "        reward = 0.5\n",
    "    elif content == \"no\":\n",
    "        reward = 0.75\n",
    "    elif content == \"maybe\":\n",
    "        reward = 1.0\n",
    "    else:\n",
    "        reward = 0.0\n",
    "    return art.Trajectory(messages_and_choices=[*messages, choice], reward=reward)\n",
    "\n",
    "\n",
    "def with_quotes(w):\n",
    "    return f\"'{w}'\"\n",
    "\n",
    "\n",
    "prompts = [\n",
    "    f\"{prefix} with {', '.join([with_quotes(w) if use_quotes else w for w in words]) if len(words) == 3 else f'{words[0]}' + (f' or {words[1]}' if len(words) > 1 else '')}\"\n",
    "    for prefix in [\"respond\", \"just respond\"]\n",
    "    for use_quotes in [True, False]\n",
    "    for words in [\n",
    "        [\"yes\", \"no\", \"maybe\"],\n",
    "        [\"maybe\", \"yes\", \"no\"],\n",
    "        [\"no\", \"yes\", \"maybe\"],\n",
    "        [\"yes\", \"maybe\", \"no\"],\n",
    "        [\"yes\", \"no\"],\n",
    "        [\"maybe\", \"no\"],\n",
    "        [\"no\", \"maybe\"],\n",
    "        [\"no\", \"yes\"],\n",
    "        [\"yes\", \"no\"],\n",
    "    ]\n",
    "]\n",
    "\n",
    "openai_client = model.openai_client()\n",
    "for _ in range(await model.get_step(), 1_000):\n",
    "    train_groups = await art.gather_trajectory_groups(\n",
    "        (\n",
    "            art.TrajectoryGroup(rollout(openai_client, prompt) for _ in range(32))\n",
    "            for prompt in prompts\n",
    "        ),\n",
    "        pbar_desc=\"gather\",\n",
    "    )\n",
    "    await model.train(\n",
    "        train_groups,\n",
    "        config=art.TrainConfig(learning_rate=1e-4),\n",
    "        # _config=art.dev.TrainConfig(\n",
    "        #     precalculate_logprobs=True,\n",
    "        # ),\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
